import numpy as np
import struct
import os
import faiss

np.set_printoptions(precision=8)


def ivecs_read(fname):
    a = np.fromfile(fname, dtype='int32')
    d = a[0]
    return a.reshape(-1, d + 1)[:, 1:].copy(), d


def fvecs_read(fname):
    data, d = ivecs_read(fname)
    return data.view('float32').astype(np.float32), d


def bvecs_read(fname):
    a = np.fromfile(fname, dtype='uint8')
    d = a[:4].view('uint8')[0]
    return a.reshape(-1, d + 4)[:, 4:].copy(), d

# store in format of vecs
def fvecs_write(filename, vecs):
    f = open(filename, "wb")
    dimension = [len(vecs[0])]

    for x in vecs:
        f.write(struct.pack('i' * len(dimension), *dimension))  # *dimension就是int, dimension就是list
        f.write(struct.pack('f' * len(x), *x))

    f.close()


def ivecs_write(filename, vecs):
    f = open(filename, "wb")
    dimension = [len(vecs[0])]

    for x in vecs:
        f.write(struct.pack('i' * len(dimension), *dimension))
        f.write(struct.pack('i' * len(x), *x))

    f.close()


def bvecs_write(filename, vecs):
    f = open(filename, "wb")
    dimension = [len(vecs[0])]

    for x in vecs:
        f.write(struct.pack('i' * len(dimension), *dimension))
        f.write(struct.pack('B' * len(x), *x))

    f.close()


def groundtruth(base, query, k):
    base_dim = base.shape[1]
    index = faiss.IndexFlatL2(base_dim)
    index.add(base)
    gnd_distance, gnd_idx = index.search(query, k)
    print("search")
    return gnd_idx


def delete_dir_if_exist(dire):
    if os.path.isdir(dire):
        # command = 'sudo rm -rf %s' % dire
        command = 'rm -rf %s' % dire
        print(command)
        os.system(command)


if __name__ == "__main__":
    old_ds_name = "glove"
    dataset_name = "glovesmall"
    k = 10
    n_query = 100
    n_base = 10000
    from_dir = '/home/zhengbian/NN_as_Classification/data/dataset/%s_50' % old_ds_name
    to_dir = '/home/zhengbian/graph-gradient/data/%s_%d' % (dataset_name, k)

    print("process %s" % dataset_name)
    query, d = fvecs_read('%s/query.fvecs' % from_dir)
    print("read query")
    # np.savetxt('query.txt', query, fmt='%.3f')

    base, d = fvecs_read('%s/base.fvecs' % from_dir)
    print("read base")
    # np.savetxt('base.txt', base, fmt='%.3f')

    query = query[:n_query, :].astype(np.float32)
    base = base[:n_base].astype(np.float32)
    gnd = groundtruth(base, query, k)
    print("base shape", base.shape, "query shape", query.shape)
    delete_dir_if_exist(to_dir)
    os.system("mkdir %s" % to_dir)
    fvecs_write('%s/query.fvecs' % to_dir, query)
    fvecs_write('%s/base.fvecs' % to_dir, base)
    ivecs_write('%s/gnd.ivecs' % to_dir, gnd)
